#include <cuda_runtime.h>
#include "conv2d.h"

extern "C" __global__ void implicit_gemm_v2(mykernelParamType param) {
    const int ohow = blockIdx.x * 16 + threadIdx.x;
    const int k = blockIdx.y * 16 + threadIdx.y;
    const int n = blockIdx.z;

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;

    __shared__ float sh_input[16][16];
    __shared__ float sh_weight[16][16];

    const int oh = ohow / param.Ow;
    const int ow = ohow % param.Ow;

    float sum = 0.0f;

    const int input_offset = n * param.c * param.h * param.w;
    const int kernel_offset = k * param.c * param.r * param.s;
    const int crs_total = param.c * param.r * param.s;

    for (int base = 0; base < crs_total; base += 16) {
        int crs_idx = base + tx;
        if (crs_idx < crs_total) {
            sh_weight[ty][tx] = param.pweight[kernel_offset + crs_idx];
        } else {
            sh_weight[ty][tx] = 0.0f;
        }

        int c = crs_idx / (param.r * param.s);
        int r = (crs_idx % (param.r * param.s)) / param.s;
        int s = crs_idx % param.s;

        int ih = oh * param.u - param.p + r;
        int iw = ow * param.v - param.q + s;

        if (ih >= 0 && ih < param.h && iw >= 0 && iw < param.w) {
            int i_addr = input_offset + c * param.h * param.w + ih * param.w + iw;
            sh_input[ty][tx] = param.pin[i_addr];
        } else {
            sh_input[ty][tx] = 0.0f;
        }

        __syncthreads();

        for (int i = 0; i < 16; i++) {
            sum += sh_input[i][tx] * sh_weight[ty][i];
        }

        __syncthreads();
    }

    if (ohow < param.Oh * param.Ow && k < param.k) {
        int o_addr = n * param.k * param.Oh * param.Ow 
                   + k * param.Oh * param.Ow 
                   + oh * param.Ow 
                   + ow;
        param.pout[o_addr] = sum;
    }
}

void launch_implicit_gemm_v2(unsigned int outh, unsigned int outw, unsigned int k, unsigned int n, mykernelParamType* param) {
    int blockx = (param->Oh * param->Ow + 15) / 16;
    int blocky = (param->k + 15) / 16;
    int blockz = param->n;
    dim3 block(16, 16, 1);
    dim3 grid(blockx, blocky, blockz);
    implicit_gemm_v2<<<grid, block>>>(*param);
}
